
import seaborn as sns
import statsmodels.api as sm

import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from sklearn.svm import SVR
from statsmodels.tsa.arima.model import ARIMA
from sklearn.metrics import mean_squared_error, accuracy_score, recall_score, precision_score, f1_score
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.patches import Rectangle
import matplotlib.cm as cm

# 1. 数据预处理和清洗
# 读取CSV文件
df = pd.read_csv('./data/数据分析/Coffee_Chain_Sales.csv')
df['Date'] = pd.to_datetime(df['Date'], format='%m/%d/%Y')
df['Year'] = df['Date'].dt.year
df['Month'] = df['Date'].dt.month
features = df[['Year', 'Month', 'AreaCode', 'Cogs', 'Margin', 'ActualandProfit', 'InventoryMargin']]
target = df['Sales']
X_train_reg, X_test_reg, y_train_reg, y_test = train_test_split(features, target, test_size=0.2, random_state=42)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train_reg)
X_test_scaled = scaler.transform(X_test_reg)
lr_model = LinearRegression()
lr_model.fit(X_train_scaled, y_train_reg)
lr_predictions = lr_model.predict(X_test_scaled)
svr_model = SVR(kernel='rbf')
svr_model.fit(X_train_scaled, y_train_reg)
svr_predictions = svr_model.predict(X_test_scaled)
X_train_ts, X_test_ts, y_train_ts, y_test_ts = train_test_split(features, target, test_size=0.2, shuffle=False)
model = ARIMA(y_train_ts, order=(1, 1, 1))
model_fit = model.fit()
yhat = model_fit.forecast(steps=len(y_test_ts))
threshold = df['Sales'].mean()
y_test_class = (y_test > threshold).astype(int)
lr_predictions_class = (lr_predictions > threshold).astype(int)
svr_predictions_class = (svr_predictions > threshold).astype(int)
arima_predictions_class = (yhat > threshold).astype(int)

lr_mse = mean_squared_error(y_test, lr_predictions_class)
lr_rmse = np.sqrt(lr_mse)
svr_mse = mean_squared_error(y_test, svr_predictions_class)
svr_rmse = np.sqrt(svr_mse)
arima_mse = mean_squared_error(y_test, arima_predictions_class)
arima_rmse = np.sqrt(arima_mse)
table_data = [
    ["Model", "Accuracy", "Recall", "Precision", "F1 Score", "MSE","RMSE"],
    ['Linear Regression', f"{accuracy_score(y_test_class, lr_predictions_class):.2f}",
     f"{recall_score(y_test_class, lr_predictions_class):.2f}",
     f"{precision_score(y_test_class, lr_predictions_class):.2f}",
     f"{f1_score(y_test_class, lr_predictions_class):.2f}",
     f"{lr_mse:.2f}",f"{lr_rmse:.2f}"],
    ['Support Vector Regression', f"{accuracy_score(y_test_class, svr_predictions_class):.2f}",
     f"{recall_score(y_test_class, svr_predictions_class):.2f}",
     f"{precision_score(y_test_class, svr_predictions_class):.2f}",
     f"{f1_score(y_test_class, svr_predictions_class):.2f}",
     f"{svr_mse:.2f}", f"{svr_rmse:.2f}"],
    ['ARIMA',f"{accuracy_score(y_test_class, arima_predictions_class):.2f}",
     f"{recall_score(y_test_class, arima_predictions_class):.2f}",
     f"{precision_score(y_test_class, arima_predictions_class):.2f}",
     f"{f1_score(y_test_class, arima_predictions_class):.2f}",
     f"{arima_mse:.2f}", f"{arima_rmse:.2f}"]
]


fig, ax = plt.subplots(figsize=(10, 7))
ax.axis('tight')
ax.axis('off')

# 设置标题和内容的渐变颜色
title_cmap = LinearSegmentedColormap.from_list("", ["blue", "lightblue"])
content_cmap = LinearSegmentedColormap.from_list("", ["grey", "white"])

# 创建表格
the_table = ax.table(cellText=table_data, loc='center', cellLoc = 'center')
the_table.auto_set_font_size(False)
the_table.set_fontsize(14)
the_table.auto_set_column_width(col=list(range(len(["Model", "Accuracy", "Recall", "Precision", "F1 Score", "MSE","RMSE"]))))

# 设置表格的列宽和行高为自适应
for (i, j), cell in the_table.get_celld().items():
    if j == 0:  # 标题行使用渐变蓝色
        cell.set_facecolor(title_cmap(0.5))
    else:
        cell.set_facecolor(content_cmap(0.5))  # 内容使用渐变白灰色
    cell.set_text_props(fontproperties=plt.matplotlib.font_manager.FontProperties(weight='bold'))

plt.savefig('model_performance_table_with_gradients.png', dpi=300, bbox_inches='tight')
plt.close()